# -*- coding: utf-8 -*-
import tensorflow as tf
from keras.callbacks import ModelCheckpoint
from keras.preprocessing import image
from keras.applications.imagenet_utils import preprocess_input, decode_predictions
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D
from keras.layers import Activation, Dropout, Flatten, Dense
from keras import backend as K
from keras import optimizers
import h5py

def train_model(train_data_dir,test_data_dir,category,
                nb_train=540,nb_test=60,img_width=224,img_height=224,
                epochs=10,batch_size = 20):
    train_data_dir = train_data_dir
    validation_data_dir = test_data_dir
    nb_train_samples = nb_train
    nb_validation_samples = nb_test
    img_width,img_height = img_width,img_height
    epochs = epochs
    batch_size = batch_size

    if K.image_data_format() == 'channels_first':
        input_shape = (3, img_width, img_height)
    else:
        input_shape = (img_width, img_height, 3)
    model = Sequential()
    model.add(Conv2D(64, (3, 3), activation='relu', padding='same',input_shape = input_shape))
    model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
    x = model.add(MaxPooling2D((2, 2), strides=(2, 2), ))

    model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
    model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
    x = model.add(MaxPooling2D((2, 2), strides=(2, 2), ))

    model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
    model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
    model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
    x = model.add(MaxPooling2D((2, 2), strides=(2, 2), ))

    model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
    model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
    model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
    x = model.add(MaxPooling2D((2, 2), strides=(2, 2), ))

    model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
    model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
    model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
    x = model.add(MaxPooling2D((2, 2), strides=(2, 2), ))
    model.load_weights('vgg16_weights_no_top.h5')
    for layer in model.layers[:25]:
        layer.trainable = False
    model.add(Flatten())
    model.add(Dense(1024))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(256))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(int(category)))
    model.add(Activation('softmax'))
    model.compile(loss='categorical_crossentropy',
                  optimizer=optimizers.SGD(lr=1e-3, momentum=0.9),
                  metrics=['accuracy'])

    train_datagen = ImageDataGenerator(
        rescale=1. / 255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True)
    savepath = "my_generator_vgg_best.h5"
    test_datagen = ImageDataGenerator(rescale=1. / 255)

    train_generator = train_datagen.flow_from_directory(
        train_data_dir,
        target_size=(img_width, img_height),
        batch_size=batch_size,
        class_mode='categorical')

    validation_generator = test_datagen.flow_from_directory(
        validation_data_dir,
        target_size=(img_width, img_height),
        batch_size=batch_size,
        class_mode='categorical')

    checkpoint = ModelCheckpoint(filepath = savepath,
                                 monitor = "val_acc",
                                 save_best_only = True,
                                 mode = "max"
            )
    callbacks_list = [checkpoint]

    model.fit_generator(
        train_generator,
        steps_per_epoch=nb_train_samples // batch_size,
        epochs=epochs,
        verbose=1,
        callbacks = callbacks_list,
        validation_data=validation_generator,
        validation_steps=nb_validation_samples // batch_size)
    model.save('webtest.h5')



